
import time
import threading
from typing import Dict, Optional

class TaskManager:
    def __init__(self):
        self.tasks: Dict[str, dict] = {}
        self.lock = threading.Lock()
    
    def create_task(self, task_id: str, task_data: dict):
        """创建新任务"""
        with self.lock:
            self.tasks[task_id] = {
                "task_id": task_id,
                "status": "pending",
                "progress": 0,
                "data": task_data,
                "created_at": time.time(),
                "video_url": None,
                "error": None
            }
    
    def update_progress(self, task_id: str, progress: int, status: str, video_url: str = None, error: str = None):
        """更新任务进度"""
        with self.lock:
            if task_id in self.tasks:
                self.tasks[task_id]["progress"] = progress
                self.tasks[task_id]["status"] = status
                if video_url:
                    self.tasks[task_id]["video_url"] = video_url
                if error:
                    self.tasks[task_id]["error"] = error
    
    def get_progress(self, task_id: str) -> Optional[dict]:
        """获取任务进度"""
        with self.lock:
            if task_id in self.tasks:
                task = self.tasks[task_id]
                return {
                    "task_id": task_id,
                    "progress": task["progress"],
                    "status": task["status"],
                    "video_url": task.get("video_url"),
                    "error": task.get("error")
                }
        return None
    
    def cleanup_old_tasks(self, max_age: int = 3600):
        """清理旧任务（默认1小时）"""
        current_time = time.time()
        with self.lock:
            expired_tasks = [
                task_id for task_id, task in self.tasks.items()
                if current_time - task["created_at"] > max_age
            ]
            for task_id in expired_tasks:
                del self.tasks[task_id]
